/*
 * Copyright (c) 2010, 2011
 *	Open Source Technologies Ltd.  All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 * 3. Neither the name of the company nor the names of its contributors
 *    may be used to endorse or promote products derived from this software
 *    without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COMPANY AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE REGENTS OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */
#include <string.h>
#include <stdlib.h>
#include <stdint.h>
#include <assert.h>

#include "defs.h"
#include "comparator.h"
#include "rbtree.h"

static rbtree_item_t *rbtree_insert_ex    __p((rbtree_t *tree, intptr_t key, intptr_t val, bool force));
static void           rbtree_insert_fixup __p((rbtree_item_t *item));
static void           rbtree_delete_fixup __p((rbtree_item_t *item));
static void           rbtree_rotate_left  __p((rbtree_item_t *item));
static void           rbtree_rotate_right __p((rbtree_item_t *item));
static rbtree_item_t *rbtree_min          __p((rbtree_item_t *item));
static rbtree_item_t *rbtree_max          __p((rbtree_item_t *item));
static rbtree_item_t *rbtree_item_create  __p((rbtree_t *tree, intptr_t key, intptr_t val));
static void           rbtree_item_replace __p((rbtree_item_t *oitem, rbtree_item_t *nitem));
static rbtree_item_t *rbtree_item_destroy __p((rbtree_item_t *item));

rbtree_t *
rbtree_create(comparator_t cmp)
{
	rbtree_t *tree;

	tree = (rbtree_t *)malloc(sizeof(rbtree_t));
	if (tree == NULL)
		return NULL;

	tree->t_magic = RBTREE_DEAD;
	rbtree_initialize(tree, cmp);

	return tree;
}

void
rbtree_initialize(rbtree_t *tree, comparator_t cmp)
{
	assert(is_rbtree_dead(tree));

	memset(tree, 0, sizeof(rbtree_t));

	tree->t_magic      = RBTREE_MAGIC;
	tree->t_comparator = cmp;
}

#ifndef NDEBUG

size_t
rbtree_size(rbtree_t *tree)
{
	assert(is_rbtree(tree));

	return tree->t_count;
}

#endif /* !NDEBUG */

rbtree_item_t *
rbtree_insert(rbtree_t *tree, intptr_t key, intptr_t val)
{
	assert(is_rbtree(tree));

	return rbtree_insert_ex(tree, key, val, false);
}

rbtree_item_t *
rbtree_replace(rbtree_t *tree, intptr_t key, intptr_t val)
{
	assert(is_rbtree(tree));

	return rbtree_insert_ex(tree, key, val, true);
}

rbtree_item_t *
rbtree_find(rbtree_t *tree, intptr_t key)
{
	rbtree_item_t *item;
	int            result;

	item = tree->t_root;
	while (item != NULL) {
		result = tree->t_comparator(key, item->ti_key);

		if (result == 0)
			break;

		if (result < 0)
			item = item->ti_left;
		else
			item = item->ti_right;
	}

	return item;
}

rbtree_item_t *
rbtree_get_first(rbtree_t *tree)
{
	assert(is_rbtree(tree));

	if (tree->t_root == NULL)
		return NULL;

	return rbtree_min(tree->t_root);
}

rbtree_item_t *
rbtree_get_last(rbtree_t *tree)
{
	assert(is_rbtree(tree));

	if (tree->t_root == NULL)
		return NULL;

	return rbtree_max(tree->t_root);
}

#ifndef NDEBUG

intptr_t
rbtree_item_get_key(rbtree_item_t *item)
{
	assert(is_rbtree_item(item));

	return item->ti_key;
}

void
rbtree_item_set_value(rbtree_item_t *item, intptr_t val)
{
	assert(is_rbtree_item(item));

	item->ti_value = val;
}

intptr_t
rbtree_item_get_value(rbtree_item_t *item)
{
	assert(is_rbtree_item(item));

	return item->ti_value;
}

#endif /* !NDEBUG */

rbtree_item_t *
rbtree_item_get_prev(rbtree_item_t *item)
{
	rbtree_t      *tree;
	rbtree_item_t *parent;

	assert(is_rbtree_item(item));

	if (item->ti_left)
		return rbtree_max(item->ti_left);

	tree = item->ti_tree;
	for (parent = item->ti_parent; parent != NULL; parent = parent->ti_parent) {
		if (tree->t_comparator(item->ti_key, parent->ti_key) > 0)
			break;
	}

	return parent;
}

rbtree_item_t *
rbtree_item_get_next(rbtree_item_t *item)
{
	rbtree_t      *tree;
	rbtree_item_t *parent;

	assert(is_rbtree_item(item));

	if (item->ti_right)
		return rbtree_min(item->ti_right);

	tree = item->ti_tree;
	for (parent = item->ti_parent; parent != NULL; parent = parent->ti_parent) {
		if (tree->t_comparator(item->ti_key, parent->ti_key) < 0)
			break;
	}

	return parent;
}

rbtree_item_t *
rbtree_item_delete(rbtree_item_t *item)
{
	rbtree_item_t *next;
	rbtree_item_t *temp;
	rbtree_item_t *pred;
	rbtree_item_t *child;

	assert(is_rbtree_item(item));

	next = rbtree_item_get_next(item);
	item->ti_tree->t_count--;

	if ((item->ti_left != NULL) && (item->ti_right != NULL)) {
		pred = rbtree_max(item->ti_left);

		temp            = pred->ti_parent;
		pred->ti_parent = item->ti_parent;
		pred->ti_left   = item->ti_left;
		pred->ti_right  = item->ti_right;
		item->ti_parent = temp;
	}

	if (item->ti_right == NULL)
		child = item->ti_left;
	else
		child = item->ti_right;

	if (item->ti_color == tc_black) {
		if (child != NULL)
			item->ti_color = child->ti_color;
		else
			item->ti_color = tc_black;

		rbtree_delete_fixup(item);
	}

	rbtree_item_replace(item, child);

	if ((item->ti_parent == NULL) && (child != NULL))
		child->ti_color = tc_black;

	rbtree_item_destroy(item);

	return next;
}

void
rbtree_finalize(rbtree_t *tree)
{
	rbtree_item_t *item;
	rbtree_item_t *tmp;

	assert(is_rbtree(tree));

	item = tree->t_root;
	while (item) {
		if (item->ti_left != NULL) {
			item = item->ti_left;
			continue;
		}
		if (item->ti_right != NULL) {
			item = item->ti_right;
			continue;
		}

		if (item->ti_parent) {
			if (item == item->ti_parent->ti_left)
				item->ti_parent->ti_left = NULL;
			else
				item->ti_parent->ti_right = NULL;
		}

		tmp  = item;
		item = item->ti_parent;

		rbtree_item_destroy(tmp);
	}
}

rbtree_t *
rbtree_destroy(rbtree_t *tree)
{
	assert(is_rbtree(tree));

	rbtree_finalize(tree);
	free(tree);

	return NULL;
}

static rbtree_item_t *
rbtree_insert_ex(rbtree_t *tree, intptr_t key, intptr_t val, bool force)
{
	rbtree_item_t *item;
	rbtree_item_t *temp;
	int            result;

	assert(is_rbtree(tree));

	item = rbtree_item_create(tree, key, val);
	if (item == NULL)
		return NULL;

	if (tree->t_root == NULL) {
		item->ti_color = tc_black;
		tree->t_root   = item;
		tree->t_count++;

		return item;
	}

	temp = tree->t_root;
	while (1) {
		result = tree->t_comparator(item->ti_key, temp->ti_key);

		if (result == 0) {
			rbtree_item_destroy(item);

			if (!force)
				return NULL;
			else
				return temp;
		}

		if (result < 0) {
			if (temp->ti_left == NULL) {
				temp->ti_left = item;
				break;
			}

			temp = temp->ti_left;
		} else {
			if (temp->ti_right == NULL) {
				temp->ti_right = item;
				break;
			}

			temp = temp->ti_right;
		}
	}

	item->ti_parent = temp;
	tree->t_count++;

	rbtree_insert_fixup(item);

	return item;
}

static void
rbtree_insert_fixup(rbtree_item_t *item)
{
	rbtree_t      *tree;
	rbtree_item_t *temp;

	assert(is_rbtree_item(item));

	tree = item->ti_tree;
	while ((item != tree->t_root) && (item->ti_parent->ti_color == tc_red)) {
		if (item->ti_parent == item->ti_parent->ti_parent->ti_left) {
			temp = item->ti_parent->ti_parent->ti_right;

			if ((temp != NULL) && (temp->ti_color == tc_red)) {
				item->ti_parent->ti_color            = tc_black;
				item->ti_parent->ti_parent->ti_color = tc_red;
				temp->ti_color                       = tc_black;

				item = item->ti_parent->ti_parent;
			} else {
				if (item == item->ti_parent->ti_right) {
					item = item->ti_parent;
					rbtree_rotate_left(item);
				}

				item->ti_parent->ti_color = tc_black;
				item->ti_parent->ti_parent->ti_color = tc_red;
				rbtree_rotate_right(item->ti_parent->ti_parent);
			}
		} else {
			temp = item->ti_parent->ti_parent->ti_left;

			if ((temp != NULL) && (temp->ti_color == tc_red)) {
				item->ti_parent->ti_color            = tc_black;
				item->ti_parent->ti_parent->ti_color = tc_red;
				temp->ti_color                       = tc_black;

				item = item->ti_parent->ti_parent;
			} else {
				if (item == item->ti_parent->ti_left) {
					item = item->ti_parent;
					rbtree_rotate_right(item);
				}

				item->ti_parent->ti_color = tc_black;
				item->ti_parent->ti_parent->ti_color = tc_red;

				rbtree_rotate_left(item->ti_parent->ti_parent);
			}
		}
	}

	tree->t_root->ti_color = tc_black;
}

static void
rbtree_delete_fixup(rbtree_item_t *item)
{
	rbtree_t      *tree;
	rbtree_item_t *w;

	assert(is_rbtree_item(item));

	tree = item->ti_tree;
	while ((item != tree->t_root) && (item->ti_color == tc_black)) {
		if (item == item->ti_parent->ti_left) {
			w = item->ti_parent->ti_right;

			if (w->ti_color == tc_red) {
				w->ti_color               = tc_black;
				item->ti_parent->ti_color = tc_red;

				rbtree_rotate_left(item->ti_parent);

				w = item->ti_parent->ti_right;
			}

			if (((w->ti_left == NULL) || (w->ti_left->ti_color == tc_black)) &&
			    ((w->ti_right == NULL) || (w->ti_right->ti_color == tc_black))) {
				w->ti_color = tc_red;
				item        = item->ti_parent;
			} else {
				if ((w->ti_right == NULL) || (w->ti_right->ti_color == tc_black)) {
					w->ti_left->ti_color = tc_black;
					w->ti_color          = tc_red;

					rbtree_rotate_right(w);

					w = item->ti_parent->ti_right;
				}

				w->ti_color               = item->ti_parent->ti_color;
				item->ti_parent->ti_color = tc_black;
				w->ti_right->ti_color     = tc_black;

				rbtree_rotate_left(item->ti_parent);

				item = tree->t_root;
			}
		} else {
			w = item->ti_parent->ti_left;

			if (w->ti_color == tc_red) {
				w->ti_color               = tc_black;
				item->ti_parent->ti_color = tc_red;

				rbtree_rotate_right(item->ti_parent);

				w = item->ti_parent->ti_left;
			}

			if (((w->ti_left == NULL) || (w->ti_left->ti_color == tc_black)) &&
			    ((w->ti_right == NULL) || (w->ti_right->ti_color == tc_black))) {
				w->ti_color = tc_red;
				item         = item->ti_parent;
			} else {
				if ((w->ti_left == NULL) || (w->ti_left->ti_color == tc_black)) {
					w->ti_right->ti_color = tc_black;
					w->ti_color           = tc_red;

					rbtree_rotate_left(w);

					w = item->ti_parent->ti_left;
				}

				w->ti_color               = item->ti_parent->ti_color;
				item->ti_parent->ti_color = tc_black;
				w->ti_left->ti_color      = tc_black;

				rbtree_rotate_right(item->ti_parent);

				item = tree->t_root;
			}
		}
	}

	item->ti_color = tc_black;
}

static void
rbtree_rotate_left(rbtree_item_t *item)
{
	rbtree_t      *tree;
	rbtree_item_t *temp;

	assert(is_rbtree_item(item));

	temp           = item->ti_right;
	item->ti_right = temp->ti_left;

	if (temp->ti_left != NULL)
		temp->ti_left->ti_parent = item;
	temp->ti_parent = item->ti_parent;

	tree = item->ti_tree;
	if (item == tree->t_root)
		tree->t_root = temp;
	else if (item == item->ti_parent->ti_left)
		item->ti_parent->ti_left = temp;
	else
		item->ti_parent->ti_right = temp;

	temp->ti_left   = item;
	item->ti_parent = temp;
}

static void
rbtree_rotate_right(rbtree_item_t *item)
{
	rbtree_t      *tree;
	rbtree_item_t *temp;

	assert(is_rbtree_item(item));

	temp          = item->ti_left;
	item->ti_left = temp->ti_right;

	if (temp->ti_right != NULL)
		temp->ti_right->ti_parent = item;
	temp->ti_parent = item->ti_parent;

	tree = item->ti_tree;
	if (item == tree->t_root)
		tree->t_root = temp;
	else if (item == item->ti_parent->ti_right)
		item->ti_parent->ti_right = temp;
	else
		item->ti_parent->ti_left = temp;

	temp->ti_right  = item;
	item->ti_parent = temp;
}

static rbtree_item_t *
rbtree_min(rbtree_item_t *item)
{
	assert(is_rbtree_item(item));

	while (item->ti_left != NULL)
		item = item->ti_left;

	return item;
}

static rbtree_item_t *
rbtree_max(rbtree_item_t *item)
{
	assert(is_rbtree_item(item));

	while (item->ti_right != NULL)
		item = item->ti_right;

	return item;
}

static rbtree_item_t *
rbtree_item_create(rbtree_t *tree, intptr_t key, intptr_t val)
{
	rbtree_item_t *item;

	assert(is_rbtree(tree));

	item = (rbtree_item_t *)malloc(sizeof(rbtree_item_t));
	if (item == NULL)
		return NULL;
	memset(item, 0, sizeof(rbtree_item_t));

	item->ti_magic = RBTREE_ITEM_MAGIC;
	item->ti_tree  = tree;
	item->ti_color = tc_red;
	item->ti_key   = key;
	item->ti_value = val;

	return item;
}

static void
rbtree_item_replace(rbtree_item_t *oitem, rbtree_item_t *nitem)
{
	rbtree_t *tree;

	assert(is_rbtree_item(oitem));
	assert(is_rbtree_item_opt(nitem));

	tree = oitem->ti_tree;
	if (oitem->ti_parent == NULL) {
		tree->t_root = nitem;
	} else {
		if (oitem == oitem->ti_parent->ti_left)
			oitem->ti_parent->ti_left = nitem;
		else
			oitem->ti_parent->ti_right = nitem;
	}

	if (nitem != NULL)
		nitem->ti_parent = oitem->ti_parent;
}

static rbtree_item_t *
rbtree_item_destroy(rbtree_item_t *item)
{
	assert(is_rbtree_item(item));

	item->ti_magic = RBTREE_ITEM_DEAD;
	free(item);

	return NULL;
}
